#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Date    : 2018-01-23 11:22:44

'''
需求：
1、主机分组
2、主机信息配置文件用configparser解析
3、可批量执行命令、发送文件，结果实时返回，执行格式如下
batch_run  -h h1,h2,h3   -g web_clusters,db_servers    -cmd  "df -h"　
batch_scp   -h h1,h2,h3   -g web_clusters,db_servers -action put
-local test.py  -remote /tmp/
4、主机用户名密码、端口可以不同
5、执行远程命令使用paramiko模块
6、批量命令需使用multiprocessing并发
'''


import paramiko
try:
    import ConfigParser
except ImportError:
    import configparser as ConfigParser
import argparse
import sys
import multiprocessing


q = multiprocessing.JoinableQueue()


class MyProcess(multiprocessing.Process):
    def __init__(self, q):
        multiprocessing.Process.__init__(self)
        self.q = q

    def scpCallBack(self, recv, send):
        pass

    def run(self):
        while True:
            ip, username, password, port, cmd = self.q.get().strip().split(',')
            if cmd.startswith('put'):
                t = paramiko.Transport((ip, int(port)))
                t.connect(username=username, password=password)
                sftp = paramiko.SFTPClient.from_transport(t)
                c, local, remote = cmd.strip().split(' ')
                try:
                    if str(sftp.stat(remote)).strip().startswith('d'):
                        remote = '/%s/%s' % (remote.strip('/'),
                                             local.split('/')[-1])
                        sftp.put(local, remote)
                        error_info = ''
                        out_info = '上传文件成功'
                    else:
                        error_info = '文件已存在'
                        out_info = ''
                except IOError:
                    try:
                        sftp.put(local, remote)
                        error_info = ''
                        out_info = '上传文件成功'
                    except:
                        error_info = '上传文件失败'
                        out_info = ''
                finally:
                    t.close()
            else:
                client = paramiko.SSHClient()
                client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
                client.connect(hostname=ip, port=port,
                               username=username, password=password)
                stdin, stdout, stderr = client.exec_command(cmd)
                error_info = stderr.read()
                out_info = stdout.read()
            self.q.task_done()
            print('%s 执行 %s 命令结果为:\n%s %s' %
                  (ip, cmd, str(out_info), str(error_info)))
            return ''


def checkOption(src, dest):
    dest_set = set(dest)
    src_set = set(src)
    if src_set.issubset(dest_set):
        return True
    else:
        return False


def getHostInfo(cfg, h):
    host_info = '%s,%s,%s,%s' % (cfg.get(h, 'ip'), cfg.get(
        h, 'username'), cfg.get(h, 'password'), cfg.get(h, 'port'))
    return host_info


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('-host', action='store', dest='host', help='目标主机')
    parser.add_argument('-g', action='store', dest='group', help='主机组')
    parser.add_argument('-cmd', action='store', dest='command', help='执行命令')
    parser.add_argument('-action', action='store', dest='action', help='执行动作')
    parser.add_argument('-local', action='store', dest='local', help='要上传的文件')
    parser.add_argument('-remote', action='store', dest='remote', help='目标路径')
    opt = parser.parse_args(sys.argv[1:])
    if opt.host is None or opt.group is None:
        print('-host 和 -g 是必须的选项，用法请看 -h')
        sys.exit(1)
    elif (opt.action is not None or opt.local is not None or opt.remote is not None) and opt.command is not None:
        print('-cmd和-action,-local,-remote不能连用')
        sys.exit(1)
    else:
        host_list = opt.host.strip().split(',')
        group_list = opt.group.strip().split(',')
        server_info_file = 'server-info.ini'
        # server_info = open(server_info_file, 'rw')
        cfg = ConfigParser.ConfigParser()
        cfg.read(server_info_file)
        # print(cfg.sections())
        if not checkOption(host_list, cfg.sections()):
            print('有不支持的主机，支持的主机列表为:%s' % str(cfg.sections().remove('group')))
            sys.exit(1)
        elif not checkOption(group_list, cfg.options('group')):
            print('有不支持的主机组，支持的主机组列表为:%s' % str(cfg.options))
            sys.exit(1)
        else:
            all_host_list = set()
            group_list_host = []
            for h in group_list:
                group_list_host.extend(cfg.get('group', h).strip().split(','))

            # 这里也可以将 host_list 和 group_list_host 变成俩个集合
            # 接着求这两个集合的交集即可

            # print('%s %s' % (str(host_list), str(group_list_host)))
            # print('%s' % str(host_list.extend(group_list_host)))
            host_list.extend(group_list_host)

            for hs in host_list:
                all_host_list.add(hs)

            print(str(all_host_list))

            prcoess_list = []
            for i in range(len(all_host_list)):
                pro = MyProcess(q)
                pro.daemon = True
                pro.start()
                prcoess_list.append(pro)
                # pro.join(timeout=60)

            # 将条目放入queue中，ip,username,password,port,cmd
            print('开始put条目进入queue中')
            if opt.command is not None:
                cmd = opt.command.strip()
            else:
                cmd = 'put %s %s' % (opt.local, opt.remote)
            for i in all_host_list:
                item = '%s,%s' % (getHostInfo(cfg, i), cmd)
                q.put(item)
            q.join()

            print('开始join进程')
            for j in prcoess_list:
                j.join()

            print('任务完成')
